# Mathmatic—Modeling—competition

# Q1：

# ```
# -*- coding: utf-8 -*-
"""
Created on Sun Sep 21 16:09:25 2025

@author: JIAMIN
"""
"""
MAT文件数据处理脚本
在Spyder中运行，用于解析.mat文件、重采样信号并生成包含元数据和处理后信号的宽表，以及展开宽表
"""

import re
import time
from pathlib import Path
import numpy as np
import pandas as pd
import scipy.io as sio
from scipy.signal import resample, welch, stft
from scipy.stats import skew, kurtosis
import pywt
from IPython.display import display  # 用于在Spyder中友好显示DataFrame

def parse_status_size_load_position_from_path(path_str: str):  #从路径/文件名解析元信息：把样本的标签(故障类型status、故障尺寸、故障位置、载荷)从文件路径和文件名自动解析出来
    """
    从文件路径/文件名中解析元数据
    - status: B / IR / OR / N
    - fault_size_inch: 故障尺寸（英寸）
    - position: 仅OR有，3->Orthogonal, 6->Centered, 12->Opposite
    - load: 0/1/2/3
    """
    p = Path(path_str)
    name = p.stem  # 文件名(无扩展)
    up = name.upper()
    parts_upper = [s.upper() for s in p.parts]  #把路径的每一段（目录名与文件名）都转成大写、放进列表，用于从目录结构中兜底查找信息。

    # 1) 解析状态(status)
    status = None
    # 先看文件名前缀
    if up.startswith("IR"):
        status = "IR"
    elif up.startswith("OR"):
        status = "OR"
    elif up.startswith("B"):
        status = "B"
    elif up.startswith("N") or "NORMAL" in up:
        status = "N"
    # 再从目录兜底
    if status is None:
        if any(seg in {"IR", "INNER", "INNER_RACE"} for seg in parts_upper):
            status = "IR"
        elif any(seg in {"OR", "OUTER", "OUTER_RACE"} for seg in parts_upper):
            status = "OR"
        elif any(seg in {"B", "BALL"} for seg in parts_upper):
            status = "B"
        elif any("NORMAL" in seg or seg == "N" for seg in parts_upper):
            status = "N"

    # 2) 解析故障尺寸(fault_size_inch)
    fault_size_inch = None
    m_size = re.search(r'(?:B|IR|OR)(\d{3})', up)
    if m_size:
        try:
            fault_size_inch = int(m_size.group(1)) / 1000.0  #文件名为007而体重要求是0.007，所以进行转换
        except (ValueError, TypeError):
            fault_size_inch = None
    elif status == "N":
        fault_size_inch = 0.0

    # 3) 解析位置(position) - 仅OR状态  #只有OR：外圈故障时才会从各个不同的方向获取数据，因为OR是固定的
    position = None
    if status == "OR":
        m_pos = re.search(r'OR\d{3}[@_-]?(\d{1,2})', up)
        pos_val = None
        if m_pos:
            pos_val = m_pos.group(1)
        else:
            # 从目录中查找位置信息
            candidates = {seg for seg in parts_upper if seg in {"3", "6", "12"}}
            if "12" in candidates:
                pos_val = "12"
            elif "6" in candidates:
                pos_val = "6"
            elif "3" in candidates:
                pos_val = "3"

        if pos_val in {"3", "6", "12"}:
            pos_map = {
                "3": "Orthogonal(3 o'clock)",
                "6": "Centered(6 o'clock)",
                "12": "Opposite(12 o'clock)"
            }
            position = pos_map[pos_val]

    # 4) 解析载荷(load)
    load = None
    m_load = re.search(r'_(\d)(?:\.MAT)?$', up)
    if m_load:
        try:
            load = int(m_load.group(1))
        except (ValueError, TypeError):
            load = None
    else:
        # 从目录中查找载荷信息
        for seg in reversed(parts_upper):
            if seg in {"0", "1", "2", "3"}:
                try:
                    load = int(seg)
                except (ValueError, TypeError):
                    load = None
                break

    return status, fault_size_inch, position, load


def infer_rpm_from_load(load: int) -> int:
    """根据载荷推断转速(RPM)，适用于12kHz CWRU数据集"""
    rpm_map = {0: 1797, 1: 1772, 2: 1750, 3: 1730}
    return rpm_map.get(load, None)


def extract_data_id(var_names):
    """从变量名列表中提取数据编号（如'X118'）"""
    for k in var_names:
        m = re.search(r'(X\d+)', k)
        if m:
            return m.group(1)
    return None


def load_mat_signals_to_features(mat_path: str):
    """读取.mat文件并提取信号特征和数据ID"""
    try:
        mat = sio.loadmat(mat_path)
        keys = [k for k in mat.keys() if not k.startswith("__")]  #.mat 文件加载后会自动包含一些以 __ 开头的系统键（如 __header__、__version__ 等，用于存储文件元信息，过滤掉这些信息

        features = {"DE_time": None, "FE_time": None, "BA_time": None, "RPM": None}
        data_id = extract_data_id(keys)

        for k in keys:
            upk = k.upper()
            # #匹配驱动端振动时域信号（键名以_DE_TIME结尾）
            if re.search(r'_DE_TIME$', upk):
                features["DE_time"] = mat[k].flatten()  # 转为一维数组
            # 匹配风扇端振动时域信号（键名以_FE_TIME结尾）
            elif re.search(r'_FE_TIME$', upk):
                features["FE_time"] = mat[k].flatten()
            # 匹配基座振动时域信号（键名以_BA_TIME结尾）
            elif re.search(r'_BA_TIME$', upk):
                features["BA_time"] = mat[k].flatten()
             # 匹配转速信息（键名以RPM结尾）
            elif re.search(r'RPM$', upk):
                # 统一RPM格式为(1,1)数组
                v = mat[k]
                v = np.array(v)
                if v.size == 1:
                    v = v.reshape(1, 1)
                features["RPM"] = v

        return features, data_id
    except Exception as e:
        print(f"[错误] 读取文件 {mat_path} 失败: {str(e)}")
        return {"DE_time": None, "FE_time": None, "BA_time": None, "RPM": None}, None


#用于将原始信号重采样并调整长度，使其统一为指定的采样率（32kHz）和时长（8 秒）
def resample_to_8s_32k(x, fs_in, target_fs=32000, target_dur=8.0):  #x：输入的原始信号（通常是一维数组）    fs_in：原始信号的采样率（单位：Hz）
    """重采样信号到目标采样率和时长"""
    if x is None:
        return None
    x = np.array(x).reshape(-1)  # x 转换为 numpy 数组，并通过 reshape(-1) 确保其为一维数组（扁平化处理）

    target_len = int(target_fs * target_dur)   #计算目标信号的总长度（采样点数）：  采样点数 = 采样率（Hz）× 时长（秒）
    n_samples = int(len(x) * target_fs / fs_in)  #按比例将原始信号长度转换为目标采样率下的长度（不考虑时长限制）
    #例如：原始信号采样率 12kHz，长度 120000 点（10 秒），目标 32kHz 时，临时长度 = 120000 × 32000 / 12000 = 320000 点

    # 重采样   #重采样需要先按比例转换长度，才能保证信号的频率特性不变
    x_resampled = resample(x, n_samples)

    # 对齐长度
    #如果重采样后的信号长度超过目标长度（256000 点），截取前 target_len 个点
    if len(x_resampled) > target_len:
        return x_resampled[:target_len]
    #如果重采样后的信号长度不足目标长度，用 0 填充（补零）到目标长度
    elif len(x_resampled) < target_len:
        return np.pad(x_resampled, (0, target_len - len(x_resampled)))
    else:
        #如果长度刚好等于目标长度，直接返回重采样后的信号
        return x_resampled



#根据文件路径或信号长度推断（推断）出信号的采样率（单位：Hz）
def infer_fs(row):
    """推断采样率"""
    path = str(row["file"]).lower()
    if "12khz" in path:
        return 12000
    if "48khz" in path:
        return 48000

    # fallback: 用长度推断
    for sig in ["DE_time", "FE_time", "BA_time"]:
        arr = row.get(sig, None)
        if arr is not None:
            n = len(arr)
            if 90000 <= n <= 110000:
                return 12000
            if 350000 <= n <= 410000:
                return 48000
    return None


#信号标准化处理
def process_source_wide(wide_df):
    processed = []
    for _, row in wide_df.iterrows():
        fs_in = infer_fs(row)  #调用infer_fs函数推断采样率
        if fs_in is None:
            print(f"[WARN] 仍无法推断采样率: {row['file']}")
            continue

        #对三种信号（DE_time/FE_time/BA_time）进行重采样l
        de_new = resample_to_8s_32k(row.get("DE_time"), fs_in)
        fe_new = resample_to_8s_32k(row.get("FE_time"), fs_in)
        ba_new = resample_to_8s_32k(row.get("BA_time"), fs_in)

        # 将处理后的结果（含元信息和新信号）添加到列表
        processed.append({
            "file": row["file"],  # 原始文件路径
            "status": row["status"],  # 状态信息
            "fault_size_inch": row["fault_size_inch"],  # 故障尺寸
            "position": row["position"],  # 位置信息
            "load": row["load"],  # 载荷信息
            "data_id": row["data_id"],  # 数据ID
            "DE_time": de_new,  # 重采样后的驱动端信号
            "FE_time": fe_new,  # 重采样后的风扇端信号
            "BA_time": ba_new,  # 重采样后的基座信号
            "RPM": row["RPM"],  # 转速信息
        })

    return pd.DataFrame(processed, columns=wide_df.columns)


def process_mat_files(root_dir=None):
    """主函数：处理指定目录下的所有.mat文件并生成宽表，之后进行重采样等处理"""
    # 设置默认目录（Spyder当前工作目录）
    if root_dir is None:
        root_dir = Path.cwd().joinpath(r"E:\研究生\数学建模\2025研究生题\数据集\源域数据集")
    else:
        root_dir = Path(root_dir)

    print(f"开始处理目录: {root_dir}")
    start_time = time.time()

    # 查找所有.mat文件
    mat_files = list(root_dir.rglob("*.mat"))
    total_files = len(mat_files)

    if total_files == 0:
        print("未找到任何.mat文件！")
        return None

    print(f"找到 {total_files} 个.mat文件，开始处理...\n")

    # 处理每个文件生成初始宽表
    rows = []
    for i, fpath in enumerate(mat_files, 1):
        # 显示进度
        if i % 5 == 0 or i == total_files:
            print(f"处理进度: {i}/{total_files} ({i/total_files:.1%})")

        # 解析元数据
        status, fault_size_inch, position, load = parse_status_size_load_position_from_path(str(fpath))

        # 读取信号特征
        feats, data_id = load_mat_signals_to_features(str(fpath))

        # 补充RPM（如果缺失）
        if feats["RPM"] is None and load is not None:
            rpm_val = infer_rpm_from_load(load)
            if rpm_val is not None:
                feats["RPM"] = np.array([[rpm_val]])


            row = {
                "file": str(fpath.relative_to(root_dir)),  # 存储相对路径（便于迁移）
                "status": status,  # 从路径解析的状态
                "fault_size_inch": fault_size_inch,  # 从路径解析的故障尺寸
                "position": position,  # 从路径解析的位置
                "load": load,  # 从路径解析的载荷
                "data_id": data_id,  # 从文件中提取的数据ID
                "DE_time": feats["DE_time"],  # 驱动端信号
                "FE_time": feats["FE_time"],  # 风扇端信号
                "BA_time": feats["BA_time"],  # 基座信号
                "RPM": feats["RPM"]  # 转速（原始或推断的）
            }
            rows.append(row)  # 将当前行数据添加到列表

    # 创建初始宽表并保存为pickle
    wide_df = pd.DataFrame(rows, columns=[
        "file", "status", "fault_size_inch", "position", "load", "data_id",
        "DE_time", "FE_time", "BA_time", "RPM"
    ])
    # 保存初始宽表为pickle文件
    wide_df.to_pickle("E:\\研究生\\数学建模\\2025研究生题\\新\\wide_df.pkl")
    print("初始宽表已保存为: wide_df.pkl")

    # 对初始宽表进行重采样等处理
    new_wide_df = process_source_wide(wide_df)

    # 计算耗时
    elapsed_time = time.time() - start_time
    print(f"\n处理完成！耗时: {elapsed_time:.2f}秒")
    if new_wide_df is not None:
        print(f"处理后宽表维度: {new_wide_df.shape}")
        if not new_wide_df.empty:
            print(f"第一行DE_time形状: {new_wide_df.iloc[0]['DE_time'].shape}")
    else:
        print("处理后宽表为空")

    return new_wide_df

def _to_1d(arr):   #统一信号为一维数组
    """把 wide_df 的 cell 转为 1D numpy 数组"""
    if arr is None:
        return None
    a = np.array(arr)
    if a.size == 0:
        return None
    if a.ndim == 1:
        return a
    if a.ndim == 2 and 1 in a.shape:
        return a.reshape(-1)
    return a.reshape(-1)

def _rpm_from_cell(cell):#提取转速单值
    """提取 RPM 单值"""
    if cell is None:
        return None
    a = np.array(cell).astype("float64", copy=False)
    if a.size == 0:
        return None
    return float(a.reshape(-1)[0])

def expand_wide_df(wide_df):#宽表转长表
    all_parts = []
    for _, row in wide_df.iterrows():
        meta = {
            "file": row["file"],          # 文件路径
            "status": row["status"],      # 状态（如正常/故障）
            "fault_size_inch": row["fault_size_inch"],  # 故障尺寸
            "position": row["position"],  # 位置（如3点/6点）
            "load": row["load"],          # 载荷
            "data_id": row["data_id"]     # 数据ID
        }

        # 提取三路信号：用_to_1d转为一维数组
        de = _to_1d(row["DE_time"])
        fe = _to_1d(row["FE_time"])
        ba = _to_1d(row["BA_time"])
        rpm_val = _rpm_from_cell(row["RPM"])

        #确定当前文件的“最大信号长度”（用于统一三路信号长度）
        max_len = max(len(x) if x is not None else 0 for x in [de, fe, ba])
        if max_len == 0:
            continue


       #对齐信号长度（补NaN或截断）
        def align(arr):
            if arr is None:
                return np.full(max_len, np.nan)
            if len(arr) < max_len:
                return np.concatenate([arr, np.full(max_len - len(arr), np.nan)])
            return arr[:max_len]

        part = pd.DataFrame({
            **meta,
            "DE_time": align(de),
            "FE_time": align(fe),
            "BA_time": align(ba),
            "RPM": np.full(max_len, rpm_val if rpm_val is not None else np.nan)
        })
        all_parts.append(part)

    #合并所有文件的长表片段，重置索引（避免重复）
    long_table = pd.concat(all_parts, ignore_index=True)
    return long_table

def time_features(x):
    """时域统计特征"""
    x = np.asarray(x).ravel()
    if len(x) == 0:
        return {}
    rms = np.sqrt(np.mean(x**2))  #均方根
    mean_abs = np.mean(np.abs(x))   #平均绝对值
    sqr_mean = np.mean(np.sqrt(np.abs(x)))  #方均根的均值
    peak = np.max(np.abs(x))   #峰值

    feats = {
        "mean": np.mean(x),
        "std": np.std(x),
        "var": np.var(x),
        "rms": rms,
        "peak": peak,
        "p2p": np.ptp(x),
        "mean_abs": mean_abs,
        "shape_factor": rms / (mean_abs + 1e-12),
        "crest_factor": peak / (rms + 1e-12),
        "impulse_factor": peak / (mean_abs + 1e-12),
        "margin_factor": peak / (sqr_mean**2 + 1e-12),
        "clearance_factor": peak / (sqr_mean + 1e-12),
        "kurtosis": kurtosis(x),
        "skewness": skew(x),
    }
    return feats

def freq_features(x, fs, fr):#频域特征（转速归一化）
    """频域特征（转速归一化）
     # x: 输入时域信号（一维数组）
    # fs: 信号采样率（Hz）
    # fr: 参考转速（通常是电机转频，Hz，即每秒转数）"""
    #f：输出的频率轴（单位：Hz），即功率谱对应的频率点；
    # Pxx：功率谱密度值，表示信号在对应频率点的功率强度；
    #nperseg=2048：将信号分成 2048 点的段进行 FFT（快速傅里叶变换），平衡频率分辨率和计算效率（段越长，频率分辨率越高）。
    f, Pxx = welch(x, fs=fs, nperseg=2048)
    Pxx = Pxx / np.sum(Pxx)  # 归一化功率谱

    # 转频归一化
    f_norm = f / fr

    #构建基础频域特征字典
    feats = {
        "spec_centroid": np.sum(f_norm * Pxx),
        "spec_bandwidth": np.sqrt(np.sum(((f_norm - np.mean(f_norm))**2) * Pxx)),
        "spec_skewness": skew(Pxx),
        "spec_kurtosis": kurtosis(Pxx),
        "spec_entropy": -np.sum(Pxx * np.log(Pxx + 1e-12)),
    }

    # 带通能量（按 fr 倍频划分）
    bands = [(0.8, 1.2), (1.8, 2.2), (2.8, 3.2), (4.5, 5.5)] #(0.8, 1.2)：接近 1 倍转频

    #遍历每个频段，用布尔掩码（mask）筛选出归一化频率在该频段内的功率谱点，求和得到该频段的能量占比，作为特征存入字典。
    for i, (lo, hi) in enumerate(bands, 1):
        mask = (f_norm >= lo) & (f_norm <= hi)
        feats[f"band_energy_{i}"] = np.sum(Pxx[mask])
    return feats


#时频域特征（STFT + 小波能量）
def tf_features(x, fs, fr):
    """时频域特征（STFT + 小波能量）"""
    # --- STFT ---
    #STFT 的核心思想：将长信号分成多个短时窗口，对每个窗口做 FFT，得到 “时间 - 频率” 二维平面的幅值分布，可直观看到不同时间点的频率成分
    f, t, Zxx = stft(x, fs=fs, nperseg=1024)
    power = np.abs(Zxx)**2
    power = power / (np.sum(power) + 1e-12)

    # 归一化频率
    f_norm = f / fr
    feats = {
        "tf_entropy": -np.sum(power * np.log(power + 1e-12)),
        "tf_mean_freq": np.sum(np.mean(power, axis=1) * f_norm),
    }

    # --- 小波分解 ---
    #小波分解的优势：相比 STFT，小波具有 “多分辨率分析” 能力 —— 低频段时间分辨率低、频率分辨率高，高频段时间分辨率高、频率分辨率低，更适合捕捉故障冲击这类 “短时高频” 信号。
    coeffs = pywt.wavedec(x, 'db4', level=4)#选择 Daubechies 4 小波（常用的小波基，适合捕捉突变信号如故障冲击）；
    energy = np.array([np.sum(c**2) for c in coeffs])
    energy_ratio = energy / (np.sum(energy) + 1e-12)
    for i, e in enumerate(energy_ratio):
        feats[f"wavelet_energy_{i}"] = e
    return feats


#从长表提取特征
#从长表按文件分组→提取每组信号的时域 / 频域 / 时频域特征→生成特征矩阵和标签→保存结果
def extract_features_from_long(long_table, signal_col="DE_time", fs=32000):
    """
    从展开后的长表中提取信号特征，并生成特征矩阵和标签序列
    核心作用：将长表中按“采样点”存储的分散数据，按“文件”聚合后提取完整特征，为建模做准备

    参数说明：
        long_table: pd.DataFrame - 展开后的长表，每行对应一个信号采样点，需包含以下关键列：
            "file": 文件名（用于按文件分组）
            "status": 信号标签（如正常/故障类型，作为输出y）
            signal_col: 指定的信号列（默认"DE_time"，即驱动端时域信号）
            "RPM": 转速信息（用于频域特征的转速归一化）
        signal_col: str - 要提取特征的信号列名，默认提取驱动端信号"DE_time"
        fs: int - 信号采样率（单位：Hz），默认32000Hz（与前期重采样标准化保持一致）

    返回值：
        X: pd.DataFrame - 特征矩阵，每行对应一个文件的所有特征，索引为文件名
        y: pd.Series - 标签序列，索引与X一致，值为每个文件的状态（如故障类型）
    """
    # 初始化列表：存储每个文件的特征字典和对应标签
    feature_rows = []  # 每个元素是一个文件的特征字典（含时域/频域/时频域特征）
    labels = []  # 每个元素是一个文件的标签（与feature_rows一一对应）

    # 1. 按文件名分组：将长表中同一文件的所有采样点聚合为一组
    # fid: 分组键（即文件名），group: 该文件对应的所有采样点数据（DataFrame）
    for fid, group in long_table.groupby("file"):
        # 2. 提取当前文件的目标信号（如DE_time），并预处理
        # group[signal_col]: 取该文件所有采样点的信号值
        # .dropna(): 删除信号中的缺失值（NaN），避免特征计算报错
        # .values: 转换为NumPy数组，适配后续特征提取函数（需数组输入）
        x = group[signal_col].dropna().values

        # 跳过空信号：若当前文件的有效信号长度为0，不进行特征提取
        if len(x) == 0:
            continue

        # 3. 计算参考转速fr（单位：Hz，即每秒转数），用于频域特征的转速归一化
        # group["RPM"].iloc[0]: 取该文件的第一个转速值（同一文件转速一致，长表中重复存储）
        rpm = group["RPM"].iloc[0]
        # 转速转换：RPM（转/分钟）→ Hz（转/秒），若转速无效（0/负数/None），默认fr=1.0避免除零错误
        fr = (rpm / 60.0) if rpm and rpm > 0 else 1.0

        # 4. 提取三类特征，并合并为一个特征字典
        feats = {}  # 存储当前文件的所有特征
        feats.update(time_features(x))  # 加入时域特征（如均值、RMS、峰度等）
        feats.update(freq_features(x, fs, fr))  # 加入频域特征（如谱重心、谱熵、带通能量等）
        feats.update(tf_features(x, fs, fr))  # 加入时频域特征（如时频熵、小波能量占比等）

        # 5. 为特征字典添加文件名（便于后续追溯文件与特征的对应关系）
        feats["file"] = fid
        # 将当前文件的特征和标签加入列表
        feature_rows.append(feats)
        labels.append(group["status"].iloc[0])  # 取该文件的第一个状态标签（同一文件标签一致）

    # 6. 构建特征矩阵X：将特征字典列表转换为DataFrame，以文件名作为索引
    X = pd.DataFrame(feature_rows).set_index("file")
    # 7. 构建标签序列y：将标签列表转换为Series，索引与X保持一致（确保特征与标签一一对应）
    y = pd.Series(labels, index=X.index, name="status")

    # 返回特征矩阵和标签序列，供后续建模（如分类、回归）使用
    return X, y

if __name__ == "__main__":
    # 在Spyder中运行时自动执行
    new_wide_df = process_mat_files()

    # 展开宽表
    if new_wide_df is not None and not new_wide_df.empty:
        long_table = expand_wide_df(new_wide_df)
        print("\n=== 展开后长表示例 ===")
        display(long_table.head())  # 显示前5行

        # 保存展开后的长表为CSV
        long_table.to_csv("E:\\研究生\\数学建模\\2025研究生题\\新\\expanded_long_table.csv", index=None)#源域数据32khz整理
        print("\n展开后长表已保存为: expanded_long_table.csv")

        # 提取特征
        X, y = extract_features_from_long(long_table, signal_col="DE_time", fs=32000)
        print("\n=== 特征提取结果 ===")
        print(f"特征矩阵形状: {X.shape}, 标签形状: {y.shape}")
        display(X.head())

        # 保存特征和标签
        X.to_csv('E:\\研究生\\数学建模\\2025研究生题\\新\\features.csv')
        y.to_csv('E:\\研究生\\数学建模\\2025研究生题\\新\\labels.csv', header=True)
        print("\n特征和标签已分别保存为: features.csv 和 labels.csv")
    else:
        print("无有效数据进行宽表展开和特征提取")

    # 保持控制台打开（Spyder中有用）
    input("\n按回车键退出...")

# ```

